#ifndef _KDTREE_H_
#define _KDTREE_H_

#include <malloc.h>
#include <cstring>
#include <cstdio>
#include <cstdlib>

namespace math 
{
template <typename value_type, unsigned dim > struct kdnode  ;
template <typename value_type, unsigned dim > struct kdhyperrect ;
template <typename value_, unsigned dim> class kdtree ;

template <typename value_type, unsigned dim>
struct res_node 
{
	typedef kdnode <value_type, dim> node_type ;
	node_type *item;
	value_type dist_sq;
	struct res_node *next;

	typedef res_node res_node_type ;
/* inserts the item. if dist_sq is >= 0, then do an ordered insert */
	static int rlist_insert(res_node_type *list, node_type *item, value_type dist_sq)
	{
		res_node_type *rnode;

		if(!(rnode = (res_node_type *) malloc (sizeof (res_node_type)))) {
			return -1;
		}
		rnode->item = item;
		rnode->dist_sq = dist_sq;

		if(dist_sq >= 0.0) {
			while(list->next && list->next->dist_sq < dist_sq) {
				list = list->next;
			}
		}
		rnode->next = list->next;
		list->next = rnode;
		return 0;
	}
};


/**
 * Result Class 
 */
template <typename value_, unsigned dim>
struct kdres 
{
	typedef value_ value_type ; 
	typedef kdres  res_type ;
	typedef res_node <value_type, dim> res_node_type ;
	typedef vector <value_type, dim > point_type ; 
	kdtree<value_type, dim> *tree;
	res_node_type *rlist, *riter;
	unsigned size;
        
	unsigned kd_res_size() const { return (this->size);}
			
	void kd_res_free()
	{
		this->clear_results();
		free(this->rlist);
//			free(rset);
	}

	void kd_res_rewind() { this->riter = this->rlist->next; }
	int kd_res_end(){ return this->riter == 0;}
	int kd_res_next(){
		this->riter = this->riter->next;
		return this->riter != 0;
	}
	void *kd_res_item_data()
	{
		return kd_res_item(this, 0);
	}

	void clear_results()
	{
		res_node_type *tmp, *node = this->rlist->next;

		while(node) {
			tmp = node;
			node = node->next;
			free(tmp);
		}

		this->rlist->next = 0;
	}

	void * kd_res_item(point_type  & pos)
	{
		if(this->riter) {
			for (unsigned i = 0; i < dim ; ++i ) pos [i] = this->riter->item->pos[i] ;

			return this->riter->item->data;
		}
		return 0;
	}
};


template <typename value_type, unsigned dim>
struct kdnode {
	typedef kdnode <value_type, dim> node_type ;
	typedef res_node <value_type, dim> res_node_type ;
	typedef kdhyperrect<value_type, dim> hyperrect_type ;
	typedef vector <value_type , dim>  point_type ;
	//value_type *pos;
	point_type pos ;
	int dir;
	void *data;
	node_type *left, *right;	/* negative/positive side */

/* Find one of the nearest nodes from the specified point.
 *
 * This function returns a pointer to a result set with at most one element.
 */
	static int find_nearest(node_type *node , const point_type & pos, value_type range, res_node_type *list, int ordered, int ndim)
	{
		value_type dist_sq, dx;
		int i, ret, added_res = 0;
		if(!node) return 0;
		dist_sq = 0;
		for(i=0; i< ndim; i++) {
			dist_sq += (node->pos[i] - pos[i]) * (node->pos[i] - pos[i]);
		}
		if(dist_sq <= (range * range)) {
			if(list->rlist_insert(list, node, ordered ? dist_sq : -1.0) == -1) {
				return -1;
			}
			added_res = 1;
		}

		dx = pos[node->dir] - node->pos[node->dir];

		ret = find_nearest(dx <= 0.0 ? node->left : node->right, pos, range, list, ordered, ndim);
		if(ret >= 0 && std::abs(dx) < range) {
			added_res += ret;
			ret = find_nearest(dx <= 0.0 ? node->right : node->left, pos, range, list, ordered, ndim);
		}
		if(ret == -1) {
			return -1;
		}
		added_res += ret;

		return added_res;
	}
	// find k_knn // Wrong implementation. 
	static int find_nearest(node_type *node , const point_type & pos, value_type range, res_node_type *list, int ordered, int ndim, int & count, int k)
	{
		if (count >= k ) return 1 ;

		value_type dist_sq, dx;
		int i, ret, added_res = 0;
		if(!node) return 0;
		dist_sq = 0;
		for(i=0; i< ndim; i++) {
			dist_sq += (node->pos[i] - pos[i]) * (node->pos[i] - pos[i]);
		}
		if(dist_sq <= (range * range)) {
			if(list->rlist_insert(list, node, ordered ? dist_sq : -1.0) == -1) {
				return -1;
			}
			count ++ ;
			added_res = 1;
		}

		dx = pos[node->dir] - node->pos[node->dir];

		ret = find_nearest(dx <= 0.0 ? node->left : node->right, pos, range,  list, ordered, ndim, count , k );
		if(ret >= 0 && std::abs(dx) < range) {
			added_res += ret;
			ret = find_nearest(dx <= 0.0 ? node->right : node->left, pos,  range , list, ordered, ndim, count , k);
		}
		if(ret == -1) {
			return -1;
		}
		added_res += ret;

		return added_res;
	}

};

/* ---- hyperrectangle helpers ---- */
template <typename value_type, unsigned dim>
struct kdhyperrect {
	
	typedef kdhyperrect<value_type, dim> hyperrect_type ;
	typedef kdnode <value_type, dim> node_type ;
	typedef vector <value_type, dim> point_type ;
	point_type min, max; 

	unsigned get_dim () {return dim ;}

	hyperrect_type* hyperrect_create(const point_type & min, const point_type & max)
	{
		size_t size = dim * sizeof(value_type);
		hyperrect_type* rect = 0;

		if (!(rect = new hyperrect_type )) {return 0 ;}
		rect->min = min ; 
		rect->max = max ; 
		return rect;
	}


	hyperrect_type* hyperrect_duplicate(const hyperrect_type *rect)
	{
		return hyperrect_create(rect->min, rect->max);
	}

	void hyperrect_extend(const point_type & pos)
	{
		this->min = math::min (this->min, pos); 
		this->max = math::max (this->max, pos); 
	}

	value_type hyperrect_dist_sq(const point_type & pos)
	{
		int i;
		value_type result = 0;

		for (i=0; i < dim; i++) {
			if (pos[i] < this->min[i]) {
				result += (this->min[i] - pos[i]) * (this->min[i] - pos[i]);
			} else if (pos[i] > this->max[i]) {
				result += (this->max[i] - pos[i]) * (this->max[i] - pos[i]);
			}
		}

		return result;
	}



	int insert_rec(node_type **nptr, const point_type & pos, void *data, int dir, int ndim)
	{
		int new_dir;
		node_type *node;

		if(!*nptr) {
			if (! (node = new node_type)) {return -1; }
			node->pos = pos ;
			node->data = data;
			node->dir = dir;
			node->left = node->right = 0;
			*nptr = node;
			return 0;
		}

		node = *nptr;
		new_dir = (node->dir + 1) % ndim;
		if(pos[node->dir] < node->pos[node->dir]) {
			return insert_rec(&(*nptr)->left, pos, data, new_dir, ndim);
		}
		return insert_rec(&(*nptr)->right, pos, data, new_dir, ndim);
	}

	void clear_rec(node_type *node, void (*destr)(void*))
	{
		if(!node) return;

		clear_rec(node->left, destr);
		clear_rec(node->right, destr);
	
		if(destr) {
			destr(node->data);
		}
		delete node ;
	}
};



/**
 * Define a kd-tree for "dim"-dimensional data 
 */
template <typename value_, unsigned dim>
	class kdtree {
public :

	typedef value_ value_type ;
	typedef kdres<value_type, dim> res_type ;
	typedef kdhyperrect <value_type, dim> hyperrect_type ;
	typedef kdnode <value_type, dim> node_type ;
	typedef res_node <value_type, dim> res_node_type ;
	typedef kdtree tree_type ;

	typedef vector <value_type, dim> point_type ; 
	node_type *root;
	hyperrect_type *rect;
	void (*destr)(void*);
//	tree_type * tree; 

public :
	unsigned get_dim () const { return dim ;}
	tree_type * get_tree () {return this ;}
public :

	kdtree()
	{
		this->root = 0;
		this->destr = 0;
		this->rect = 0;
	}

    /* free the struct kdtree */
	~kdtree()
	{
		kd_clear();
	}


    /* remove all the elements from the tree */
	void kd_clear()
	{
		rect->clear_rec(this->root, this->destr);
		this->root = 0;

		if (this->rect) {
			//free(rect);
			delete rect ;
			this->rect = 0;
		}
	}

	/**
	 * Add elements. 
	 */
	int kd_insert( const point_type & pos, void *data)
	{
		if (this->rect == 0) {
			this->rect = rect->hyperrect_create(pos, pos);
		} else {
			rect->hyperrect_extend(pos);
		}

		if (rect->insert_rec(&this->root, pos, data, 0, dim)) {
			return -1;
		}
		return 0;
	}


    /* if called with non-null 2nd argument, the function provided
	 * will be called on data pointers (see kd_insert) when nodes
	 * are to be removed from the tree.
	 */

	void kd_data_destructor(struct kdtree *tree, void (*destr)(void*))
	{
		tree->destr = destr;
	}

	void kd_nearest_i(node_type *node, const point_type &pos, node_type **result, value_type *result_dist_sq, hyperrect_type* rect)
	{
		int dir = node->dir;
		int i, side;
		value_type dummy, dist_sq;
		node_type *nearer_subtree, *farther_subtree;
		value_type *nearer_hyperrect_coord, *farther_hyperrect_coord;

		/* Decide whether to go left or right in the tree */
		dummy = pos[dir] - node->pos[dir];
		if (dummy <= 0) {
			nearer_subtree = node->left;
			farther_subtree = node->right;
			nearer_hyperrect_coord = & rect->max [dir];
			farther_hyperrect_coord = & rect->min [dir];
			side = 0;
		} else {
			nearer_subtree = node->right;
			farther_subtree = node->left;
			nearer_hyperrect_coord = & rect->min [dir];
			farther_hyperrect_coord = & rect->max [dir];
			side = 1;
		}

		if (nearer_subtree) {
			/* Slice the hyperrect to get the hyperrect of the nearer subtree */
			dummy = *nearer_hyperrect_coord;
			*nearer_hyperrect_coord = node->pos[dir];
			/* Recurse down into nearer subtree */
			kd_nearest_i(nearer_subtree, pos, result, result_dist_sq, rect);
			/* Undo the slice */
			*nearer_hyperrect_coord = dummy;
		}

		/* Check the distance of the point at the current node, compare it
		 * with our best so far */
		dist_sq = 0;
		for(i=0; i < dim; i++) {
			dist_sq += (node->pos[i] - pos[i]) * (node->pos[i] - pos[i]);
		}
		if (dist_sq < *result_dist_sq) {
			*result = node;
			*result_dist_sq = dist_sq;
		}

		if (farther_subtree) {
			/* Get the hyperrect of the farther subtree */
			dummy = *farther_hyperrect_coord;
			*farther_hyperrect_coord = node->pos[dir];
			/* Check if we have to recurse down by calculating the closest
			 * point of the hyperrect and see if it's closer than our
			 * minimum distance in result_dist_sq. */
			if (rect->hyperrect_dist_sq(pos) < *result_dist_sq) {
				/* Recurse down into farther subtree */
				kd_nearest_i(farther_subtree, pos, result, result_dist_sq, rect);
			}
			/* Undo the slice on the hyperrect */
			*farther_hyperrect_coord = dummy;
		}
	}

	void * kd_nearest(const point_type & pos, point_type & res_pos)
	{
		res_type * presult = kd_nearest_impl (pos ); 
		void * data =  presult->kd_res_item(res_pos );
		presult->kd_res_free(); 
		free (presult);
		return data; 
	}

	res_type * kd_nearest_impl(const point_type &pos)
	{
		hyperrect_type *rect;
		node_type *result;
		res_type *rset;
		value_type dist_sq;
		int i;

		if (!this) return 0;
		if (!this->rect) return 0;

		/* Allocate result set */
		if(!(rset =  (res_type *)malloc(sizeof *rset))) {
			return 0;
		}

		if(!(rset->rlist = (res_node_type *) malloc(sizeof(res_node_type)))) {
			free(rset);
			return 0;
		}
		rset->rlist->next = 0;
		rset->tree = this;

		/* Duplicate the bounding hyperrectangle, we will work on the copy */
		if (!(rect = rect->hyperrect_duplicate(this->rect))) {
			rset->kd_res_free();
			return 0;
		}

		/* Our first guesstimate is the root node */
		result = this->root;
		dist_sq = 0;
		for (i = 0; i < dim; i++)
			dist_sq += (result->pos[i] - pos[i]) * (result->pos[i] - pos[i]);

		/* Search for the nearest neighbour recursively */
		kd_nearest_i(this->root, pos, &result, &dist_sq, rect);

		/* Free the copy of the hyperrect */
		free(rect);

		/* Store the result */
		if (result) {
			if (rset->rlist->rlist_insert(rset->rlist, result, -1.0) == -1) {
				rset->kd_res_free();
				return 0;
			}
			rset->size = 1;
			rset->kd_res_rewind();
			return rset;
		} else {
			rset->kd_res_free();
			return 0;
		}
	}

	res_type * kd_nearest_range_n(const point_type & pos, value_type range ,int k, unsigned ordered = 0)
	{
		int ret;
		res_type *rset;
		int count = 0;

		if(!(rset = (res_type *)malloc(sizeof *rset))) {
			return 0;
		}
		if(!(rset->rlist = (res_node_type *)malloc (sizeof(res_node_type)))) {
			free(rset);
			return 0;
		}
		rset->rlist->next = 0;
		rset->tree = this ;

		if((ret = root->find_nearest(this->root , pos, range, rset->rlist, ordered, dim, count, k)) == -1) {
			rset->kd_res_free();
			free(rset) ;
			return 0;
		}
		rset->size = ret;
		rset->kd_res_rewind();
		return rset;

	}

	res_type * kd_nearest_range(const point_type & pos, value_type range, unsigned ordered = 0)
	{
		int ret;
		res_type *rset;

		if(!(rset = (res_type *)malloc(sizeof *rset))) {
			return 0;
		}
		if(!(rset->rlist = (res_node_type *)malloc (sizeof(res_node_type)))) {
			free(rset);
			return 0;
		}
		rset->rlist->next = 0;
		rset->tree = this ;

		if((ret = root->find_nearest(this->root, pos, range, rset->rlist, ordered, dim)) == -1) {
			rset->kd_res_free();
			free(rset) ;
			return 0;
		}
		rset->size = ret;
		rset->kd_res_rewind();
		return rset;
	}

};
}

#endif	/* _KDTREE_H_ */
